package saramax

import (
	"context"
	"encoding/json"
	"github.com/IBM/sarama"
	"go.uber.org/zap"
	"time"
)

type BatchHandler[T any] struct {
	fn            func(msg []*sarama.ConsumerMessage, t []T) error
	batchSize     int
	batchDuration time.Duration
}

//type BatchHandlerOption func(*BatchHandler[receiver T])

func NewBatchHandler[T any](fn func(msg []*sarama.ConsumerMessage, t []T) error) *BatchHandler[T] {
	batchHandler := &BatchHandler[T]{
		fn:            fn,
		batchSize:     10,
		batchDuration: time.Second * 5,
	}
	return batchHandler
}

//	func WithBatchSize(batchSize int) BatchHandlerOption {
//		return func(b *BatchHandler[any]) {
//			b.batchSize = batchSize
//		}
//	}
//
//	func WithBatchDuration(batchDuration time.Duration) BatchHandlerOption {
//		return func(b *BatchHandler[any]) {
//			b.batchDuration = batchDuration
//		}
//	}
func (b BatchHandler[T]) Setup(session sarama.ConsumerGroupSession) error {
	return nil
}

func (b BatchHandler[T]) Cleanup(session sarama.ConsumerGroupSession) error {
	return nil
}

// 批量消费
func (b BatchHandler[T]) ConsumeClaim(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
	ch := claim.Messages()
	for {
		msgs := make([]*sarama.ConsumerMessage, 0, b.batchSize)
		ts := make([]T, 0, b.batchSize)
		ctx, cancel := context.WithTimeout(context.Background(), b.batchDuration)
		defer cancel()

		for i := 0; i < b.batchSize; i++ {
			select {
			case <-ctx.Done():
				//超时了，立即返回
				break
			case msg, ok := <-ch:
				//管道关闭了
				if !ok {
					return nil
				}

				var t T
				err := json.Unmarshal(msg.Value, &t)
				if err != nil {
					zap.L().Error("反序列化消息体失败",
						zap.String("topic", msg.Topic),
						zap.Int32("partition", msg.Partition),
						zap.Int64("offset", msg.Offset),
						// 这里也可以考虑打印 msg.Value，但是有些时候 msg 本身也包含敏感数据
						zap.Error(err))
					session.MarkMessage(msg, "")
					continue
				}
				ts = append(ts, t)
				msgs = append(msgs, msg)
			}
		}
		if err := b.fn(msgs, ts); err == nil {
			// 这边就要都提交了
			for _, msg := range msgs {
				session.MarkMessage(msg, "")
			}
		} else {
			// 这里可以考虑重试，也可以在具体的业务逻辑里面重试
			// 也就是 eg.Go 里面重试
		}
		cancel()
	}
}
